package com.slipper.ai.service;

import com.google.gson.reflect.TypeToken;
import com.slipper.ai.flow.FlowEngine;
import com.slipper.ai.flow.Parser;
import com.slipper.ai.flow.vo.FlowVO;
import com.slipper.ai.flow.vo.MessageVO;
import com.slipper.ai.flow.vo.NodeVO;
import com.slipper.ai.flow.vo.ResultVO;
import com.slipper.ai.model.chat.ChatModelFactory;
import com.slipper.ai.prompt.Prompts;
import com.slipper.ai.vector.KnowledgeVectorService;
import com.slipper.constants.Constants;
import com.slipper.core.cache.CacheUtil;
import com.slipper.exception.CommonException;
import com.slipper.system.ai.bo.AiApplicationBO;
import com.slipper.system.ai.bo.AiFlowBO;
import com.slipper.system.ai.bo.AiModelBO;
import com.slipper.system.ai.service.AiFlowService;
import com.slipper.system.ai.service.AiModelService;
import com.slipper.util.BeanUtil;
import com.slipper.util.ResultUtil;
import com.slipper.util.Tool;
import jakarta.annotation.Resource;
import org.springframework.ai.chat.messages.AssistantMessage;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.chat.messages.SystemMessage;
import org.springframework.ai.chat.messages.UserMessage;
import org.springframework.ai.chat.model.ChatModel;
import org.springframework.ai.chat.prompt.Prompt;
import org.springframework.stereotype.Component;

import java.lang.reflect.Type;
import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;
import java.util.stream.Collectors;

@Component
public class ChatService {
    @Resource
    private AiModelService aiModelService;
    @Resource
    private FlowEngine flowEngine;
    @Resource
    private AiFlowService aiFlowService;
    @Resource
    private KnowledgeVectorService knowledgeVectorService;

    public String generatePrompt(String userMessage) {
        if (Tool.isBlank(userMessage)) throw new CommonException("用户消息不能为空！");

        AiModelBO aiModel = aiModelService.queryDefaultChatModel();
        ChatModel chatModel = ChatModelFactory.create(aiModel);
        if (chatModel == null) throw new CommonException("未配置默认大语言模型！");

        Prompt prompt = new Prompt()
                .augmentSystemMessage(Prompts.GENERATE_LLM_PROMPT)
                .augmentUserMessage(userMessage);
        return chatModel.call(prompt)
                .getResults()
                .stream()
                .map(item -> item.getOutput().getText())
                .collect(Collectors.joining("\n\n"));
    }

    public String chat(String systemMessage, String userMessage, AiModelBO aiModel) {
        return chat(null, null, systemMessage, userMessage, aiModel);
    }

    public String chat(
            String sessionId,
            Integer historyCount,
            String systemMessage,
            String userMessage,
            AiModelBO aiModel) {
        if (Tool.isBlank(userMessage)) throw new CommonException("用户消息不能为空！");

        ChatModel chatModel = ChatModelFactory.create(aiModel);
        if (chatModel == null) throw new CommonException("未配置默认大语言模型！");

        List<Message> messageList = new ArrayList<>();
        if (Tool.isNotBlank(systemMessage)) {
            messageList.add(new SystemMessage(systemMessage));
        }

        List<MessageVO> historyList = null;
        String HISTORY_PREFIX = "HISTORY_PREFIX_";
        if (historyCount != null && historyCount > 0 && Tool.isNotBlank(sessionId)) {
            String cacheKey = HISTORY_PREFIX + sessionId;
            Type type = new TypeToken<List<MessageVO>>(){}.getType();
            historyList = CacheUtil.get(cacheKey, type);
            if (historyList != null) {
                historyList.forEach(message -> {
                    if (Tool.isNotBlank(message.getUserMessage())) {
                        messageList.add(new UserMessage(message.getUserMessage()));
                    }
                    if (Tool.isNotBlank(message.getAssistantMessage())) {
                        messageList.add(new AssistantMessage(message.getAssistantMessage()));
                    }
                });
            }
        }

        messageList.add(new UserMessage(userMessage));
        String assistantMessage = chatModel.call(new Prompt(messageList))
                .getResults()
                .stream()
                .map(item -> item.getOutput().getText())
                .collect(Collectors.joining("\n\n"));

        if (historyCount != null && historyCount > 0 && Tool.isNotBlank(sessionId)) {
            if (historyList == null) {
                historyList = new ArrayList<>();
            }
            historyList.add(new MessageVO(userMessage, assistantMessage));

            while(historyList.size() > historyCount) {
                historyList.removeFirst();
            }

            String cacheKey = HISTORY_PREFIX + sessionId;
            CacheUtil.set(cacheKey, historyList, 30 * 60);
        }

        return assistantMessage;
    }

    public String chat(AiApplicationBO bo, boolean debug) {
        String userMessage = bo.getQuery();
        if (Tool.isBlank(userMessage)) return ResultUtil.failureResult("用户消息不能为空！");
        AiModelBO aiModel = aiModelService.queryBy(bo.getModelId());

        if (Constants.NO.equals(bo.getUseFlow())) {
            if (Tool.isNotBlank(bo.getModelParamJson())) {
                aiModel.setParamJson(BeanUtil.concatMapJson(aiModel.getParamJson(), bo.getModelParamJson()));
            }

            String similarityText = null;
            if (Tool.isNotBlank(bo.getKnowledgeId())) {
                similarityText = knowledgeVectorService.similarityText(bo.getKnowledgeId(), userMessage);
                if (Tool.isNotBlank(similarityText)) {
                    userMessage = Prompts.RAG_PROMPT.formatted(similarityText, userMessage);
                }
            }

            String result = this.chat( bo.getSessionId(), bo.getRecordSize(), bo.getPrompt(), userMessage, aiModel);
            if (debug) {
                Map<String, String> resultMap = new HashMap<>();
                resultMap.put("knowledge", similarityText);
                resultMap.put("result", result);
                return ResultUtil.successResult("调试成功。", resultMap);
            } else {
                return ResultUtil.successResult("成功。", result);
            }
        } else {
            if (Tool.isBlank(bo.getFlowId())) return ResultUtil.failureResult("请先配置流程名称！");

            AiFlowBO flowBO = aiFlowService.queryBy(bo.getFlowId());
            if (flowBO == null) {
                return ResultUtil.failureResult("当前流程名称已废弃，请重新配置流程名称！");
            }
            if (Tool.isBlank(flowBO.getJsonText())) {
                return ResultUtil.failureResult("流程【" + flowBO.getName() + "】未配置节点信息！");
            }

            FlowVO flow = Parser.parse(flowBO.getJsonText());
            if (flow == null) return ResultUtil.failureResult("流程配置解释出错！");

            Map<String, Object> flowData = new HashMap<>();
            String finalUserMessage = userMessage;
            flow.getNodes().stream()
                    .filter(node -> NodeVO.TYPE_START.equals(node.getType()))
                    .findFirst()
                    .ifPresent(node -> node.getOutputList().forEach(output -> {
                        if ("question".equals(output.getId())) {
                            flowData.put(output.getName(), finalUserMessage);
                        }
                    }));
            flow.setFlowData(flowData);
            flow.setSessionId(bo.getSessionId());
            ResultVO result = flowEngine.start(flow);

            if (debug) {
                Map<String, String> resultMap = new HashMap<>();
                resultMap.put("result", result.getMessage() == null ? result.getContent() : result.getMessage());
                return ResultUtil.successResult("调试成功。", resultMap);
            } else {
                if (Tool.isNotBlank(result.getMessage())) {
                    return ResultUtil.failureResult("失败", result.getMessage());
                } else {
                    return ResultUtil.successResult("成功。", result.getContent());
                }
            }
        }
    }
}
